import numpy as np
import random


def iid(opt, train_label):
    """
    Sample I.I.D. client data from dataset
    """
    num_items = int(len(train_label) / opt.num_users)
    dict_users, all_idxs = {}, [i for i in range(len(train_label))]
    for i in range(opt.num_users):
        dict_users[i] = np.random.choice(all_idxs, num_items, replace=False)
        all_idxs = list(set(all_idxs) - set(dict_users[i]))
    return dict_users


def noniid(opt, train_label, n_list, k_list):
    """
    Sample non-I.I.D client data from dataset
    根据非独立同分布（Non-I.I.D.）方式划分客户端数据集。

    参数：
    - opt: 包含数据集和划分设置的选项对象。
    - train_label: 训练数据的标签（通常是独热编码）。
    - n_list: 每个客户端需要选择的类别数列表。
    - k_list: 每个客户端从每个类别选择的样本数量列表。

    返回：
    - dict_users: 每个客户端的样本索引字典。
    - classes_list: 每个客户端选择的类别列表。
    """
    # 设置数据的划分数量（shards）和每个类别的图片数量。
    num_shards, num_imgs = 0, 0
    if opt.dataset == 'FashionVC':
        # 16,862 training imgs
        num_shards, num_imgs = 27, 624 # FashionVC 数据集划分为 27 个块，每块 624 张图片。
    elif opt.dataset == 'Ssense':
        # 13,696 training imgs
        num_shards, num_imgs = 28, 489
    idx_shard = [i for i in range(num_shards)] #创建包含所有块索引的列表。
    dict_users = {} # 用于存储每个客户端的数据索引。
    idxs = np.arange(num_shards * num_imgs) #生成一个包含所有样本的索引
    labels = [np.argmax(item) for item in train_label] #将独热编码标签转换为类别标签。
    labels = np.array(labels)[:num_shards * num_imgs] #保留有效的标签部分。
    # sort labels
    idxs_labels = np.vstack((idxs, labels)) # 将样本索引和标签堆叠起来。
    idxs_labels = idxs_labels[:, idxs_labels[1, :].argsort()] # 按标签排序样本。
    idxs = idxs_labels[0, :] # 提取排序后的样本索引。

    # 记录每个类别在数据集中的起始位置。
    label_begin = {}
    cnt = 0
    for i in idxs_labels[1, :]:
        if i not in label_begin:
            label_begin[i] = cnt
        cnt += 1

    # the num for each class
    class_num = [label_begin[item+1] - label_begin[item] for item in range(opt.num_class-1)] # 计算每个类别中的样本数量。
    class_num.append(num_shards*num_imgs-label_begin[opt.num_class-1])

    class_sort = np.array(class_num)[np.array(class_num).argsort()] #将类别的样本数按升序排序。
    # each class for k_len
    # 计算每个类别可以分成多少组，每组的大小由 train_shots_max （每个客户端最多的训练样本数量）限制
    class_divided = np.floor(np.array(class_num) / opt.train_shots_max)
    idxs_class_divided = np.vstack((np.zeros(opt.num_class).astype(int), class_divided.astype(int)))

    classes_list = [] #存储每个客户端选择的类别。
    for i in range(opt.num_users):
        # 查找当前可供选择的类别
        class_candidate = [] #存储可以选择的类别，这些类别在 idxs_class_divided[1] 中的样本数量大于或等于 1。
        for index, each_class in enumerate(idxs_class_divided[1]):
            if each_class >= 1:
                class_candidate.append(index)
        # 每个客户端需要选择的类别数量（n）和每个类别的样本数（k）。
        n = n_list[i]
        k = k_list[i]
        k_len = opt.train_shots_max # 每个类别的样本数上限
        #从候选类别中随机选择 n 个类别。
        classes = random.sample(class_candidate, n)
        classes = np.sort(classes)

        print("user {:d}: {:d}-way {:d}-shot".format(i + 1, n, k))
        print("classes:", classes)
        # 为客户端选择数据
        user_data = np.array([])
        for each_class in classes:
            # 确定每个类别的数据起始位置并提取 k 个样本
            begin = idxs_class_divided[0][each_class] * k_len + label_begin[each_class.item()]
            user_data = np.concatenate((user_data, idxs[begin: begin + k]), axis=0)#用于存储每个客户端选择的数据索引。对于每个类别，按 k 的样本数量选择数据。
        dict_users[i] = user_data # 存储客户端数据索引
        classes_list.append(classes)  # 存储客户端选择的类别

        # classes里面出现了的类的数量在idxs_class_divided[1]中减1，以确保每个类别的样本不会重复选择。
        for each_class in classes:
            idxs_class_divided[0][each_class] = idxs_class_divided[0][each_class] + 1
            idxs_class_divided[1][each_class] = idxs_class_divided[1][each_class] - 1

    return dict_users, classes_list #dict_users：每个客户端的数据索引字典，存储每个客户端选择的数据。 classes_list: 存储每个客户端选择的类别列表。


def noniid_unequal(opt, train_label):
    """
    Sample non-I.I.D client data from dataset s.t clients
    have unequal amount of data
    """
    #num_shards 是数据集被划分的块（shards）数目，num_imgs 是每个块中包含的图片数量。
    num_shards, num_imgs = 0, 0
    if opt.dataset == 'FashionVC':
        num_shards, num_imgs = 843, 20  # FashionVC数据集有843个块，每块包含20张图片
    elif opt.dataset == 'Ssense':
        num_shards, num_imgs = 228, 60

    # 初始化块索引列表和用户数据字典
    idx_shard = [i for i in range(num_shards)]  # 块索引列表
    dict_users = {i: np.array([]) for i in range(opt.num_users)} # 每个用户初始分配的数据为空

    # 生成所有样本的全局索引，并提取样本标签
    idxs = np.arange(num_shards * num_imgs) # 所有样本的索引
    labels = [np.argmax(item) for item in train_label] # 提取每个样本的标签
    labels = np.array(labels)[:num_shards * num_imgs] # 截取前num_shards * num_imgs个标签


    # 对标签进行排序，以便按标签分块
    idxs_labels = np.vstack((idxs, labels)) # 将索引和标签堆叠
    idxs_labels = idxs_labels[:, idxs_labels[1, :].argsort()] # 按标签排序
    idxs = idxs_labels[0, :] # 提取排序后的索引

    # Minimum and maximum shards assigned per client:设定每个客户端可以获得的最小和最大数据块数目。每个客户端的数据块数目将从 min_shard 到 max_shard 之间随机选取。
    min_shard = 1
    max_shard = 30

    # Divide the shards into random chunks for every client 随机为每个客户端分配数据块数目
    # s.t the sum of these chunks = num_shards
    #为每个客户端随机生成一个数据块数目，保证这些块的总数等于 num_shards。这个分配是根据每个客户端数据块数量的比例来调整的，
    random_shard_size = np.random.randint(min_shard, max_shard + 1, size=opt.num_users)
    random_shard_size = np.around(random_shard_size / sum(random_shard_size) * num_shards) #这个操作可能会使得块数大于num_shards
    random_shard_size = random_shard_size.astype(int) #最终确保所有客户端的块总数加起来等于 num_shards。

    # Assign the shards randomly to each client
    # caused by np.around
    # 处理分块数目超出总块数的情况
    if sum(random_shard_size) > num_shards:
        #经过这轮校正后，能确保分块数目小于总块数
        for i in range(opt.num_users):
            # First assign each client 1 shard to ensure every client has at least one shard of data
            rand_set = set(np.random.choice(idx_shard, 1, replace=False)) # 随机选择1个块
            idx_shard = list(set(idx_shard) - rand_set) # 更新未分配的块
            for rand in rand_set:
                dict_users[i] = np.concatenate((dict_users[i], idxs[rand * num_imgs:(rand + 1) * num_imgs]), axis=0)

        random_shard_size = random_shard_size - 1 # 剩余块数减1

        # Next, randomly assign the remaining shards  将剩余的数据块随机分配给客户端
        for i in range(opt.num_users):
            if len(idx_shard) == 0:  # 若没有剩余块，则停止分配
                continue
            shard_size = random_shard_size[i] #shard_size是客户端i需要的块数
            if shard_size > len(idx_shard):  # 若客户端需要的块数超过剩余块数，则调整，这段代码确保了 shard_size 的值不超过 idx_shard 的长度（即可供选择的块数）。
                shard_size = len(idx_shard)
            rand_set = set(np.random.choice(idx_shard, shard_size, replace=False))  # 随机选择数据块，大小为这个客户端需要的块，replace代表不允许重复采样
            idx_shard = list(set(idx_shard) - rand_set)  # 更新剩余块
            for rand in rand_set:
                dict_users[i] = np.concatenate((dict_users[i], idxs[rand * num_imgs:(rand + 1) * num_imgs]), axis=0)
    else:
        # 分块数目在允许范围内的情况，直接分配
        for i in range(opt.num_users):
            shard_size = random_shard_size[i]
            rand_set = set(np.random.choice(idx_shard, shard_size, replace=False))
            idx_shard = list(set(idx_shard) - rand_set)
            for rand in rand_set:
                dict_users[i] = np.concatenate((dict_users[i], idxs[rand * num_imgs:(rand + 1) * num_imgs]), axis=0)
        # 若仍有剩余块，分配给数据最少的客户端
        if len(idx_shard) > 0:
            # Add the leftover shards to the client with minimum images:
            shard_size = len(idx_shard)
            # Add the remaining shard to the client with lowest data
            k = min(dict_users, key=lambda x: len(dict_users.get(x)))  # 找到数据最少的客户端
            rand_set = set(np.random.choice(idx_shard, shard_size, replace=False))
            for rand in rand_set:
                dict_users[k] = np.concatenate((dict_users[k], idxs[rand * num_imgs:(rand + 1) * num_imgs]), axis=0)

    return dict_users

